﻿from numpy import *
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import functools

def show_centroids(xs, cs, centroids):
    fig = plt.figure()
    ax = fig.add_subplot(111)
    plot_scatters(ax, xs, cs)
    ax.scatter(centroids[:,0], centroids[:,1], s=60.0, marker='+', color='r')
    plt.show()

def show_scatters(xs, cs, index=(0,1), axis_labels=("","")):
    fig = plt.figure()
    ax = fig.add_subplot(111)
    plot_scatters(ax, xs, cs, index, axis_labels)
    plt.show()

def plot_scatters(ax, xs, cs, index=(0,1), axis_labels=("","")):
    """显示数据的散列图
    Parameters
    ----------
    data: matrix
        显示的数据，每一行为一个样本
    labels: list[string]
        数据分类，与`matrix`的行数相同，表示样本的类型。
    index: list[int]
        选择`matrix`的两列作为散列数据
    axisLabels: list[string]
        坐标轴x，y显示的文字
    """

    color = ['r','g','y','m','b']

    labelTypes = set(cs)
    print(labelTypes)
    types = []
    typeIndex = 0
    for t in labelTypes:
        print(index)
        x = xs[:,index[0]].copy()
        y = xs[:,index[1]].copy()

        size = 0
        for i in range(len(cs)):
            if cs[i] == t:
                x[size] = xs[i,index[0]]
                y[size] = xs[i,index[1]]
                size += 1
        print(size)
        print(x[0:size])
        types.append(ax.scatter(x[0:size], y[0:size], 40.0, color[typeIndex%len(color)]))
        print(color[typeIndex%len(color)]);
        typeIndex += 1
    ax.legend(types, labelTypes)
    #ax.axis();

    #ax.xlabel(axis_labels[0])
    #ax.ylabel(axis_labels[1])


def getTreeDepth(tree):
    """获取树的深度"""
    if type(tree) == list:
        maxDepth = 0
        for key in tree[1].keys():
            depth = 1 + getTreeDepth(tree[1][key])
            if maxDepth < depth:
                maxDepth = depth
        return maxDepth
    else:
        return 1;

def getTreeWidth(tree):
    """获取tree叶子数目"""
    if type(tree) == list:
        width = 0
        for key in tree[1].keys():
            width += getTreeWidth(tree[1][key])
        return width
    else:
        return 1;

def plotNode(ax, point, text):
    ax.annotate(text, xy=point, xycoords='axes fraction',textcoords='axes fraction', va='center', ha='center', bbox=dict(boxstyle="round4"), arrowprops=dict(arrowstyle="<-"))

def plotArrow(ax, a, b):
    ax.annotate("", xy=a, xytext=b, arrowprops=dict(arrowstyle="<-"))


def plotMiddleText(ax, a, b, text):
    """在两点之间显示文字"""
    pt = [(a[i]+b[i])/2.0 for i in range(len(a))]
    ax.text(pt[0], pt[1], text, va="center", ha="center", rotation=30)

def plotTree(ax, area, tree):
    def topCenter(area):
        return ((area[0]+area[2])/2, area[3])
    #print(tree)
    (x,y) = ((area[0]+area[2])/2, area[3])
    print((x,y))
    if type(tree) == list:
        width = getTreeWidth(tree)
        depth = getTreeDepth(tree)
        plotNode(ax, (x,y), str(tree[0]))

        subarea = [area[0], area[1], 0, area[3] - (area[3]-area[1])*1.0/(depth-1)]
        for key in tree[1].keys():
            subwidth = getTreeWidth(tree[1][key])

            
            subarea[2] =  subarea[0] + (area[2]-area[0]) * 1.0 / width * subwidth;
            
            plotArrow(ax, (x,y), topCenter(subarea))
            plotMiddleText(ax, (x,y), ((subarea[0]+subarea[2])/2, subarea[3]), key)
            
            plotTree(ax, subarea, tree[1][key])
            subarea[0] = subarea[2]


    else:

        plotNode(ax, (x,y), str(tree))


    

def showTree(tree):
    """显示树型数据的图形
    Parameter
    ---------
    tree: 
    """
    fig = plt.figure(1)
    fig.clf()
    ax = plt.subplot(111, frameon=False)
    plotTree(ax, (0, 0, 1, 1), tree)
    plt.show()



def plotCurve(ax, x, expr, labels=None):
    """绘制曲线

    Parameters
    ----------
    ax:
        画布
    x: [Float]
        曲线x的值
    expr: String
        计算y点的表达式字符串
    labels: [String,String]
        x,y轴的说明，没有则采用["x", expr]
    """

    y = eval(expr)
    ax.plot(x, y)
    
    _labels = labels if labels != None else ["x", expr]
    plt.xlabel(_labels[0])
    plt.ylabel(_labels[1])


def showCurves(rowcol, xs, exprs, labels=None):
    for i in range(len(xs)):
        ax = plt.subplot(rowcol*10+i+1)
        plotCurve(ax, xs[i], exprs if type(exprs)==str else exprs[i], labels[i] if labels != None else None)

    
    plt.show()

    
def plotContour(ax, x, y, func):
    X, Y = meshgrid(x, y)
    Z = func(X, Y)
    CS = ax.contour(X,Y,Z)
    plt.clabel(CS, inline=1, fontsize=10)
    
def plotArrows(ax, points, labels):
    leafNode = dict(boxstyle="round4", fc="0.8")
    arrow_args = dict(arrowstyle="<-")

    ax.annotate('', xy=points[0], xytext=points[1], 
            va="center", ha="center", bbox=leafNode, arrowprops=arrow_args )

    ax.annotate('',  xy=points[1], xytext=points[2], 
            va="center", ha="center", bbox=leafNode, arrowprops=arrow_args )

    ax.annotate('', xy=points[2], xytext=points[3], 
            va="center", ha="center", bbox=leafNode, arrowprops=arrow_args )

    ax.text(points[0][0], points[0][1], labels[0])
    ax.text(points[1][0], points[1][1], labels[1])
    ax.text(points[2][0], points[2][1], labels[2])
    ax.text(points[3][0], points[3][1], labels[3])

